// -*- C++ -*-

#ifndef MATRIX_H
#define MATRIX_H

extern "C" {
#include <gsl/gsl_matrix.h>
#include <gsl/gsl_blas.h>
#include <gsl/gsl_permutation.h>
#include <gsl/gsl_linalg.h>
#include <gsl/gsl_eigen.h>
}

#include <iostream>
#include <vector>
#include <cassert>

class Vector;
class Matrix;

Vector operator+(const Vector&, const Vector&);
Vector operator-(const Vector&, const Vector&);
Vector operator*(const Vector&, double);
double operator*(const Vector&, const Vector&);

Vector operator*(const Matrix&, const Vector&);
Vector operator*(const Vector&, const Matrix&);

Matrix operator+(const Matrix&, const Matrix&);
Matrix operator-(const Matrix&, const Matrix&);
Matrix operator*(const Matrix&, const Matrix&);
Matrix operator*(const Matrix&, double);

class Vector {
  friend class Matrix;

  friend Vector operator+(const Vector&, const Vector&);
  friend Vector operator-(const Vector&, const Vector&);
  friend Vector operator*(const Vector&, double);
  friend double operator*(const Vector&, const Vector&);

  friend Vector operator*(const Matrix&, const Vector&);
  friend Vector operator*(const Vector&, const Matrix&);
public:
  Vector() {
    vector = 0;
  }

  Vector(int n) {
    vector = gsl_vector_alloc(n);
  }

  ~Vector() {
    if (vector) gsl_vector_free(vector);
  }

  Vector(const Vector& copy) {
    if (copy.vector) {
      vector = gsl_vector_alloc(copy.num());
      gsl_vector_memcpy(vector, copy.vector);
    } else {
      vector = 0;
    }    
  }

  Vector& operator=(const Vector& copy) {
    if (this == &copy) return *this;
    if (vector) {
      gsl_vector_free(vector);
      vector = 0;
    }
    if (copy.vector) {
      vector = gsl_vector_alloc(copy.num());
      gsl_vector_memcpy(vector, copy.vector);
    }
    return *this;
  }

  double& operator()(int i) {
    assert(i >= 0 && i < (int)vector->size);
    return vector->data[vector->stride * i];
  }

  const double& operator()(int i) const {
    assert(i >= 0 && i < (int)vector->size);
    return vector->data[vector->stride * i];
  }

  double norm() const {
    return gsl_blas_dnrm2(vector);
  }

  int num() const {
    return vector->size;
  }

  enum Exception {
    BAD_SIZE
  };
  
private:
  gsl_vector* vector;
};

inline Vector operator+(const Vector& p, const Vector& q) {
  if (p.num() != q.num()) {
    throw Vector::BAD_SIZE;
  }
  Vector r = p;
  gsl_vector_add(r.vector, q.vector);
  return r;
}

inline Vector operator-(const Vector& p, const Vector& q) {
  if (p.num() != q.num()) {
    throw Vector::BAD_SIZE;
  }
  Vector r = p;
  gsl_vector_sub(r.vector, q.vector);
  return r;
}

inline Vector operator*(const Vector& p, double q) {
  Vector r = p;
  gsl_vector_scale(r.vector, q);
  return r;
}

inline Vector operator*(double p, const Vector& q) {
  return q * p;
}

inline Vector operator/(const Vector& p, double q) {
  return p * (1.0 / q);
}

inline double operator*(const Vector& p, const Vector& q) {
  if (p.num() != q.num()) {
    throw Vector::BAD_SIZE;
  }
  double r;
  gsl_blas_ddot(p.vector, q.vector, &r);
  return r;
}

inline std::ostream& operator<<(std::ostream& os, const Vector& vector) {
  os << "[ ";
  for (int i = 0; i < vector.num(); ++i) {
    os << vector(i);
    if (i + 1 != vector.num()) os << ", ";
  }
  os << " ]";
  return os;
}

inline std::istream& operator>>(std::istream& is, Vector& vec) {
  std::vector<double> data;
  char c;
  if (!(is >> c)) {
    return is;
  }
  if (c != '[') {
    is.putback(c);
    is.setstate(std::ios_base::failbit);
    return is;
  }
  while (is >> c && c !=']') {
    is.putback(c);
    double value;
    if (!(is >> value)) {
      return is;
    }
    if (is >> c && (c != ',' && c != ';')) {
      is.putback(c);
    }
    data.push_back(value);
  }
  if (!is) {
    return is;
  }
  vec = Vector(data.size());
  for (int i = 0; i < data.size(); ++i) {
    vec(i) = data[i];
  }
  return is;
}

class Matrix {
  friend Matrix operator+(const Matrix&, const Matrix&);
  friend Matrix operator-(const Matrix&, const Matrix&);
  friend Matrix operator*(const Matrix&, const Matrix&);
  friend Matrix operator*(const Matrix&, double);

  friend Vector operator*(const Matrix&, const Vector&);
  friend Vector operator*(const Vector&, const Matrix&);
public:

  Matrix() {
    matrix = 0;
  }
  
  Matrix(int m, int n) {
    matrix = gsl_matrix_alloc(m, n);
  }

  ~Matrix() {
    if (matrix) gsl_matrix_free(matrix);
  }

  Matrix(const Matrix& copy) {
    if (copy.matrix) {
      matrix = gsl_matrix_alloc(copy.rowNum(), copy.colNum());
      gsl_matrix_memcpy(matrix, copy.matrix);
    } else {
      matrix = 0;
    }
  } 

  Matrix& operator=(const Matrix& copy) {
    if (this == &copy) return *this;
    if (matrix) {
      gsl_matrix_free(matrix);
      matrix = 0;
    }
    if (copy.matrix) {
      matrix = gsl_matrix_alloc(copy.rowNum(), copy.colNum());
      gsl_matrix_memcpy(matrix, copy.matrix);
    }
    return *this;
  }

  double& operator()(int i, int j) {
    assert(i >= 0 && i < (int)matrix->size1);
    assert(j >= 0 && j < (int)matrix->size2);
    return matrix->data[i * matrix->tda + j];
  } 

  const double& operator()(int i, int j) const {
    assert(i >= 0 && i < (int)matrix->size1);
    assert(j >= 0 && j < (int)matrix->size2);
    return matrix->data[i * matrix->tda + j];
  } 

  int rowNum() const {
    return matrix->size1;
  }

  int colNum() const {
    return matrix->size2;
  }

  double determinant() const {
    if (rowNum() != colNum()) {
      throw BAD_SIZE;
    }
    gsl_matrix* copy = gsl_matrix_alloc(rowNum(), colNum());
    gsl_permutation* perm = gsl_permutation_alloc(rowNum());
    int sig;
    gsl_matrix_memcpy(copy, matrix);
    gsl_linalg_LU_decomp(copy, perm, &sig);
    double ret = gsl_linalg_LU_det(copy, sig);
    gsl_permutation_free(perm);
    gsl_matrix_free(copy);
    return ret;
  }

  double trace() const {
    if (rowNum() != colNum()) {
      throw BAD_SIZE;
    }
    double r = 0.0;
    for (int i = 0; i < rowNum(); ++i) {
      r += operator()(i, i);
    }
    return r;
  }

  Matrix inverse() const {
    if (rowNum() != colNum()) {
      throw BAD_SIZE;
    }
    gsl_matrix* copy = gsl_matrix_alloc(rowNum(), colNum());
    gsl_permutation* perm = gsl_permutation_alloc(rowNum());
    int sig;
    gsl_matrix_memcpy(copy, matrix);
    gsl_linalg_LU_decomp(copy, perm, &sig);
    double det = gsl_linalg_LU_det(copy, sig);
    if (det == 0.0) {
      gsl_permutation_free(perm);
      gsl_matrix_free(copy);
      throw NULL_MATRIX;
    }
    Matrix r(rowNum(), colNum());
    gsl_linalg_LU_invert(copy, perm, r.matrix);
    gsl_permutation_free(perm);
    gsl_matrix_free(copy);
    return r;
  }

  Vector inverseMul(const Vector& q) const {
    if (rowNum() != colNum()) {
      throw BAD_SIZE;
    }
    Vector r(rowNum());
    gsl_matrix* copy = gsl_matrix_alloc(rowNum(), colNum());
    gsl_permutation* perm = gsl_permutation_alloc(rowNum());
    int sig;
    gsl_matrix_memcpy(copy, matrix);
    gsl_linalg_LU_decomp(copy, perm, &sig);
    int ret = gsl_linalg_LU_solve(copy, perm, q.vector, r.vector);
    gsl_permutation_free(perm);
    gsl_matrix_free(copy);
    if (ret != 0) {
      throw NULL_MATRIX;
    }
    return r;
  }

  Matrix transpose() const {
    Matrix r(colNum(), rowNum());
    gsl_matrix_transpose_memcpy(r.matrix, matrix);
    return r;
  }

  Matrix eigenVectors() const {
    if (rowNum() != colNum()) {
      throw BAD_SIZE;
    }
    gsl_matrix* tmp = gsl_matrix_alloc(rowNum(), colNum());
    gsl_matrix_memcpy(tmp, matrix);
    gsl_vector* eval = gsl_vector_alloc(rowNum());
    gsl_eigen_symmv_workspace* ws = gsl_eigen_symmv_alloc(rowNum());
    Matrix evec(rowNum(), colNum());
    gsl_eigen_symmv(tmp, eval, evec.matrix, ws);
    gsl_eigen_symmv_free(ws);
    gsl_vector_free(eval);
    gsl_matrix_free(tmp);
    return evec;
  }

  Vector eigenValues() const {
    if (rowNum() != colNum()) {
      throw BAD_SIZE;
    }
    Vector eval(rowNum());
    gsl_matrix* tmp = gsl_matrix_alloc(rowNum(), colNum());
    gsl_matrix_memcpy(tmp, matrix);
    gsl_eigen_symm_workspace* ws = gsl_eigen_symm_alloc(rowNum());
    gsl_eigen_symm(tmp, eval.vector, ws);
    gsl_eigen_symm_free(ws);
    gsl_matrix_free(tmp);
    return eval;
  }

  Matrix cholesky() {
    if (rowNum() != colNum()) {
      throw BAD_SIZE;
    }
    Matrix r(*this);
    gsl_linalg_cholesky_decomp(r.matrix);
    for (int i = 0; i < colNum(); ++i) {
      for (int j = 0; j + 1 < rowNum(); ++j) {
	r(i, j) = 0.0;
      }
    }
    return r;
  }

  Matrix perturb(double epsilon = 0.01) {
    if (rowNum() != colNum()) {
      throw BAD_SIZE;
    }
    Matrix r(*this);
    for (int i = 0; i < colNum(); ++i) {
      r(i, i) += epsilon;
    }
    return r;
  }

  enum Exception {
    BAD_SIZE, NULL_MATRIX
  };

private:

  gsl_matrix* matrix;
};

inline Matrix operator+(const Matrix& p, const Matrix& q) {
  if (p.rowNum() != q.rowNum() || p.colNum() != q.colNum()) {
    throw Matrix::BAD_SIZE;
  }
  Matrix r = p;
  gsl_matrix_add(r.matrix, q.matrix);
  return r;
}

inline Matrix operator-(const Matrix& p, const Matrix& q) {
  if (p.rowNum() != q.rowNum() || p.colNum() != q.colNum()) {
    throw Matrix::BAD_SIZE;
  }
  Matrix r = p;
  gsl_matrix_sub(r.matrix, q.matrix);
  return r;
}

inline Matrix operator*(const Matrix& p, const Matrix& q) {
  if (p.colNum() != q.rowNum()) {
    throw Matrix::BAD_SIZE;
  }
  Matrix r(p.rowNum(), q.colNum());
  gsl_blas_dgemm(CblasNoTrans, CblasNoTrans, 1.0, 
		 p.matrix, q.matrix, 0.0, r.matrix);
  return r;
}

inline Matrix operator*(const Matrix& p, double q) {
  Matrix r(p);
  gsl_matrix_scale(r.matrix, q);
  return r;
}

inline Matrix operator*(double p, const Matrix& q) {
  return q * p;
}

inline Matrix operator/(const Matrix& p, double q) {
  return p * (1.0 / q);
}

inline Vector operator*(const Matrix& p, const Vector& q) {
  if (p.colNum() != q.num()) {
    throw Matrix::BAD_SIZE;
  }
  Vector r(p.rowNum());
  gsl_blas_dgemv(CblasNoTrans, 1.0, p.matrix, q.vector, 0.0, r.vector);
  return r;
}

inline Vector operator*(const Vector& p, const Matrix& q) {
  if (q.rowNum() != p.num()) {
    throw Matrix::BAD_SIZE;
  }
  Vector r(q.colNum());
  gsl_blas_dgemv(CblasTrans, 1.0, q.matrix, p.vector, 0.0, r.vector);
  return r;
}


inline Matrix diad(const Vector& p, const Vector& q) {
  Matrix r(p.num(), q.num());
  for (int i = 0; i < p.num(); ++i) {
    for (int j = 0; j < q.num(); ++j) {
      r(i, j) = p(i) * q(j);
    }
  }
  return r;
}


inline double mulTrace(const Matrix& p, const Matrix& q) {
  if (p.rowNum() != q.colNum()) {
    throw Matrix::BAD_SIZE;
  }
  if (p.colNum() != q.rowNum()) {
    throw Matrix::BAD_SIZE;
  }
  double r = 0.0;
  for (int i = 0; i < p.rowNum(); ++i) {
    for (int j = 0; j < p.colNum(); ++j) {
      r += p(i, j) * q(j, i);
    }
  }
  return r;
}

inline std::ostream& operator<<(std::ostream& os, const Matrix& matrix) {
  os << "[ ";
  for (int i = 0; i < matrix.rowNum(); ++i) {
    os << "[ ";
    for (int j = 0; j < matrix.colNum(); ++j) {
      os << matrix(i, j);
      if (j + 1 != matrix.colNum()) os << ", ";
    }
    os << " ]";
    if (i + 1 != matrix.rowNum()) os << ", ";
  }
  os << " ]";
  return os;
}

inline std::istream& operator>>(std::istream& is, Matrix& matrix) {
  int width = 0, height = 0;
  std::vector<double> data;
  char c;
  if (!(is >> c)) {
    return is;
  }
  if (c != '[') {
    is.putback(c);
    is.setstate(std::ios_base::failbit);
    return is;
  }
  while (is >> c && c !=']') {
    if (c != '[') {
      is.putback(c);
      is.setstate(std::ios_base::failbit);
      return is;
    }
    int w = 0;    
    while (is >> c && c != ']') {
      is.putback(c);
      double value;
      if (!(is >> value)) {
        return is;
      }
      if (is >> c && (c != ',' && c != ';')) {
        is.putback(c);
      }
      ++w;
      data.push_back(value);
    }
    if (!is) {
      return is;
    }
    if (height == 0) {
      width = w;
    } 
    if (w != width){
      is.setstate(std::ios_base::failbit);
      return is;
    }
    ++height;
    if (is >> c && (c != ',' && c != ';')) {
      is.putback(c);
    }
  }
  if (!is) {
    return is;
  }
  matrix = Matrix(height, width);
  for (int i = 0; i < height; ++i) {
    for (int j = 0; j < width; ++j) {
      matrix(i, j) = data[i * width + j];
    }
  }
  return is;
}

#endif
